#include "Main.h"
#define MAX_CODE_LENGTH 16


// "Smallest possible transparent PNG" thank you for keeping me sane.
struct ZNode{
	bool last;
	int Value;
	ZNode *Child[2];

	ZNode(){last = false;Value =0;Child[0]=Child[1]=NULL;}
	ZNode(int e){last = true; Value =e;Child[0]=Child[1]=NULL;}
};

struct ZTree{
ZNode Firstnode;

ZTree(){Firstnode = ZNode();}
void CreateNodes(unsigned short int code,unsigned int value, BYTE bitlength){
	ZNode *cur = &Firstnode;
	unsigned short int bit;
	for(int i=bitlength; i>1;i--){
		bit = (code<<(16-i));
		bit >>= 15;
		if(cur->Child[bit] == NULL){cur ->Child[bit] =  new ZNode();}
		cur = cur ->Child[bit];
		}
	bit = (code&1); //last bit
	cur ->Child[bit] =  new ZNode(value);
	}

void CreateTree(int amount, BYTE maxcodelength,BYTE *CdLgnth,unsigned short int *BL_Count){
	BYTE *NC_Dict = new BYTE[maxcodelength];

	unsigned int code = 0,len=0; //next codes for all code lengths
	NC_Dict[0]=0;
    for (BYTE bits = 1; bits <= maxcodelength-1; bits++) {
		code = (code + BL_Count[bits-1]) << 1;NC_Dict[bits] = code;}

	for (int n = 0;  n <= amount-1; n++){len =  CdLgnth[n]; //assign all codes
	if (len != 0) {CreateNodes(NC_Dict[len],n,len); NC_Dict[len]++; } }
	
	}

};

class Unzipper{
//Dictionary
ZTree Codes;//Huffman Trees
ZTree DistCodes;
int Read; //Read Byte Buffer;
BYTE RBits; //Read Bit Buffer
int Write; //Output Byte buffer
//Data
public:
BYTE* in;
BYTE* out;

Unzipper(){}
void Load(BYTE* o){out =o;Write=0;}
void Input(BYTE* i){in = i;Read =0;RBits=0;}


//1.1: A DEFLATE stream is a little-endian bit stream. This explains so much.

unsigned short int RevBits(unsigned int in, BYTE amount){
	unsigned short int ret=0;
	int bit=0;
	for(int i =0; i<amount; i++){
	bit = (1<<i);
	if(in&bit){ret |= (1<<(amount-i-1));}
	}
	return ret;
}

unsigned short int GetBits(BYTE amount){ // this reads the last bits of the byte first in big endian, then reverses them.
unsigned int dat = (in[Read+2]<<16)|(in[Read+1]<<8)| (in[Read]);//16 bits can only cover up to 3 bytes.
dat >>= RBits;
dat <<= (32-amount); //get rid of bits ahead
dat >>=(32-amount);
RBits +=amount;
while(RBits >= 8){Read++;RBits-=8;}
//return RevBits(dat,amount);
return dat; //BIG RISK
}//up to 16 bits per call.

BYTE GetBit(){
BYTE ret = in[Read];
ret >>= RBits; 
ret &= 1; //if its one is one than its one.
RBits ++;
if(RBits >= 8){Read++;RBits-=8;}
return ret;
}

void SetUpDict(bool Preset){
	BYTE Code_Lengths[287];//Code length (used in docompression and creating dict)
	BYTE Dist_Lengths[32];//^
	unsigned short int BL_Count[MAX_CODE_LENGTH]={0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0}; //Bit length count for Dictionary Code lengths 0-15
	unsigned short int BL_DCount[8]={0,0,0,0,0,0,0,0}; //Bit length count for Distance Coce lengths 0-5

	if(Preset){
		for(int i=0; i<=143; i++){Code_Lengths[i] = 8;}
		for(int i=144; i<=255; i++){Code_Lengths[i] = 9;}
		for(int i=256; i<=279; i++){Code_Lengths[i] = 7;}
		for(int i=280; i<=287; i++){Code_Lengths[i] = 8;}
		for(int i=0; i<=31; i++){Dist_Lengths[i] = 5;}
		BL_Count[7]=24;
		BL_Count[8]=152;
		BL_Count[9]=112;
		BL_DCount[5]=32;
	}
	else{ //Dynamic Huffman Codes
		// So the trees are huffman encoded, so I pretty much gotta make a mini copy of this function here
		
		ZTree DictCodes;

		unsigned int HLIT = GetBits(5)+257;
		BYTE HDIST=GetBits(5)+1; //number of listed distance codes, sometimes it stops short and the rest are 0
		BYTE HCLEN=GetBits(4)+4;; //number of code codes. This will be hard to remember.
		BYTE BL_ORDER[19]={16, 17, 18,0, 8, 7, 9, 6, 10, 5, 11, 4, 12, 3, 13, 2, 14, 1, 15};
		unsigned short BL_COUNT[8] ={0,0,0,0,0,0,0,0};
		BYTE CODE_LENGTH_LENGTHS[19]={0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,};
		//the code lengths for the code lengths are 3 bit integers
		unsigned short g;
		for(int l=0;l<HCLEN;l++){
			g=GetBits(3); 
			CODE_LENGTH_LENGTHS[BL_ORDER[l]]=g;  BL_COUNT[g]++;}//assign bit lengths and count them
		BL_COUNT[0]=0;
		DictCodes.CreateTree(19,8,&CODE_LENGTH_LENGTHS[0], &BL_COUNT[0]);

		BYTE Dictout[320]; //max HDIST+HLEN
		//Dictout is the dictionary output as one bit stream, it must be read this way.
		ZNode *dc;
		unsigned int wrt=0;bool mtc=false; unsigned int cpy;
		while(wrt < (HLIT+HDIST)){ //read HLIT+HDIST codes, they have to be together at first because they may reference each other
			dc = &DictCodes.Firstnode;
			BYTE dabit;
			while(mtc == false){
				dabit = GetBits(1);
			dc = dc->Child[dabit];
			if((!dc)){MessageBoxA ( NULL,"Zlib Datastream Problem: bad dictionary dictionary.", "Something Went Wrong", MB_OK );}
			mtc = dc->last;
			}
			mtc = false;
			mtc=false;
			switch(dc->Value){
			case 16: cpy=(3+GetBits(2)); for(int i=0;i<cpy;i++){Dictout[wrt+i]=Dictout[wrt-1];}wrt+=cpy;break; //copy last byte X times
			case 17: cpy=(3+GetBits(3)); for(int i=0;i<cpy;i++){Dictout[wrt+i]=0;}wrt+=cpy;break; // copy X 0's
			case 18: cpy=(11+GetBits(7));for(int i=0;i<cpy;i++){Dictout[wrt+i]=0;}wrt+=cpy;break;//copy X 0's a lot
			default:Dictout[wrt]=dc->Value;wrt++;break;// append to stream
			}
		}
		for(int i=0; i<HLIT;i++){Code_Lengths[i]=Dictout[i];BL_Count[Dictout[i]]++;} //copy each length tree to the corresponding tree, counting bit lengths
		BL_Count[0]=0;
		for(int i=0; i<HDIST;i++){Dist_Lengths[i]=Dictout[HLIT+i];BL_DCount[Dictout[HLIT+i]]++;}
		BL_DCount[0]=0;
	}
	BL_Count[0]=0;
	BL_DCount[0]=0;
	Codes.CreateTree(287,15,Code_Lengths,BL_Count);
	DistCodes.CreateTree(32,8,Dist_Lengths,BL_DCount);
}

void Decompress(){
Read = 0;
RBits=0;
BYTE CMF = in[0],FLG= in[1]; 

short int check = ((CMF<<8) | (FLG));
check = check % 31; //if this is true then we likely got ourselves a datastream
if(check ==0){
BYTE CM,CINFO, FDICT,FLEVEL;   
CM = ((CMF)&(8));			    //00001111 Method can only be "Deflate"
CINFO =((CMF)&(240))>>4;			//11110000 Window Size (probably 7) // doesnt matter
//FCHECK =(CMF>>3);			// Checksum that makes this check above = 0
FDICT  = ((FLG)& (32))>>5;		// XX100000 //Is there a Dictionary?
FLEVEL = (FLG)& (192)>>6;	    // 11000000 "The information in FLEVEL is not needed for decompression;
									//it is there to indicate if recompression might be worthwhile."

if(CM!=8){MessageBoxA ( NULL,"Invalid PNG Datastream: Non DEFLATE compression", "Image Error", MB_OK );} //we always use DEFLATE algorithm
//int SLIDING_WINDOW = 2^(CINFO+8); kinda useless because our sliding window is the whole datastream
Read = 2; //byte read position
if(FDICT){Read =6;} //useless Dict Checksum, but we need to add 4 bytes
Inflate();
//ALDER32 CHECKSUM(4 More Bytes)
}
else{MessageBoxA ( NULL,"Invalid PNG Datastream: Check stream Fail.", "Image Error", MB_OK );}
} 

void Inflate(){
unsigned int LEN,DIST;
bool match = false;

//3 header bits
bool Last = false;
while(!Last){
	Last = (bool)GetBits(1);
    BYTE comptype = GetBits(2); 
	if(comptype ==0){ //no compression
		if(RBits !=0){Read++;}RBits =0; //disreguard any leftoever bits

		LEN = (in[Read+1]<<8)|(in[Read]);
		//ones complement of len = in[Read+2]<<8 | in[Read+3]; 
		Read +=4;
		memcpy(&out[Write], &in[Read],LEN);
		Write+=LEN;
		}
	else if(comptype !=0){
	/*
	 00 - no compression
     01 - compressed with fixed Huffman codes   (guessing 2)
     10 - compressed with dynamic Huffman codes (guessing 1)
     11 - reserved (error)
	*/
		SetUpDict((comptype == 1)? true:false); //if its 2, its preset.

		ZNode *zc;
		ZNode *dst;

		while(true){
			zc = &Codes.Firstnode;
			while(match == false){
				zc = zc->Child[GetBit()];
				if((!zc)){MessageBoxA ( NULL,"Zlib Datastream Problem:Error Decoding.", "Something Went Wrong", MB_OK );return;}
				match = zc->last;
				}
			match = false;
			if(zc->Value<256){out[Write] = zc->Value; Write++;}
			else if(zc->Value == 256){break;}
			else if(zc->Value >256){
			switch(zc->Value){ //Too lazy to figure out anything clever.
				 case 257: LEN =3; break;
				 case 258: LEN =4; break;
				 case 259: LEN =5; break;	
				 case 260: LEN =6; break;
				 case 261: LEN =7; break;
				 case 262: LEN =8; break;
				 case 263: LEN =9; break;
		 		 case 264: LEN =10; break;
				 case 265: LEN =11+ GetBits(1); break;
				 case 266: LEN =13+ GetBits(1); break;
				 case 267: LEN =15+ GetBits(1); break;
				 case 268: LEN =17+ GetBits(1); break;
				 case 269: LEN =19+ GetBits(2); break;
				 case 270: LEN =23+ GetBits(2); break;
				 case 271: LEN =27+ GetBits(2); break;
				 case 272: LEN =31+ GetBits(2); break;
				 case 273: LEN =35+ GetBits(3); break;
				 case 274: LEN =43+ GetBits(3); break;
				 case 275: LEN =51+ GetBits(3); break;
				 case 276: LEN =59+ GetBits(3); break;
				 case 277: LEN =67+ GetBits(4); break;
				 case 278: LEN =83+ GetBits(4); break;
				 case 279: LEN =99+ GetBits(4); break;
				 case 280: LEN =115+GetBits(4); break;
			 	 case 281: LEN =131+GetBits(5); break;
				 case 282: LEN =163+GetBits(5); break;
				 case 283: LEN =195+GetBits(5); break;
				 case 284: LEN =227+GetBits(5); break;
				 default: LEN =258; break;
			}

		    dst = &DistCodes.Firstnode;
		    while(match == false){
				dst = dst->Child[GetBits(1)];
				if((!dst)){MessageBoxA ( NULL,"Zlib Datastream Problem:Error Decoding for Distance.", "Something Went Wrong", MB_OK );}
				match = dst->last;
			}
	     	match = false;
			switch(dst->Value){
				case 0:   DIST=1;break;
				case 1:   DIST=2;break;
				case 2:   DIST=3;break;
				case 3:   DIST=4;break;
				case 4:   DIST=5+GetBits(1);break;
				case 5:   DIST=7+GetBits(1);break;
				case 6:   DIST=9+GetBits(2);break;
				case 7:   DIST=13+GetBits(2);break;
				case 8:   DIST=17+GetBits(3);break;
				case 9:   DIST=25+GetBits(3);break;
				case 10:  DIST=33+GetBits(4);break;
				case 11:  DIST=49+GetBits(4);break;
				case 12:  DIST=65+GetBits(5);break;
				case 13:  DIST=97+GetBits(5);break;
				case 14:  DIST=129+GetBits(6);break;
				case 15:  DIST=193+GetBits(6);break;
				case 16:  DIST=257+GetBits(7);break;
				case 17:  DIST=385+GetBits(7);break;
				case 18:  DIST=513+GetBits(8);break;
				case 19:  DIST=769+GetBits(8);break;
				case 20:  DIST=1025+GetBits(9);break;
				case 21:  DIST=1537+GetBits(9);break;
				case 22:  DIST=2049+GetBits(10);break;
				case 23:  DIST=3073+GetBits(10);break;
				case 24:  DIST=4097+GetBits(11);break;
				case 25:  DIST=6145+GetBits(11);break;
				case 26:  DIST=8193+GetBits(12);break;
				case 27:  DIST=12289+GetBits(12);break;
				case 28:  DIST=16385+GetBits(13);break;
				default: DIST = 24577+ GetBits(13);break;//seems dangerous
			}
			for(int v=0; v<LEN; v++){out[Write+v]= out[Write-DIST+v];}
			Write +=LEN;
			}
		}
	}
  }
}



};

class Zipper{
//to be made, or maybe not... I don't know.
BYTE* out;

};


